# ===========================================================================
CONF_THRESH = 0.25
IOU_THRESHOLD = 0.45
categories = ["slagcar", "car", "tricar", "motorbike", "bicycle", "bus", "truck", "tractor"]
bbuild = "build"
pt_name = 'last.pt'

model_type = "n"
model_imgsz = 768
model_e = 60
model_b = 80
model_name = "yolov5{}_s{}_e{}_b{}".format(model_type, model_imgsz, model_e, model_b)

if __name__ == "__main__":
    model_type = "s"
    model_imgsz = 704
    model_e = 100
    model_b = 64
    model_name = "local_yolov5{}_s{}_e{}_b{}".format(model_type, model_imgsz, model_e, model_b)

import logging as log

log.basicConfig(level=log.DEBUG)
# ============================================================================

import glob
import os
import random
import time
import cv2
import ctypes
import shutil
import threading
import numpy as np
from tqdm import tqdm

import struct
import torch
import json
import sys

sys.path.append("/project/train/src_repo/yolov5_js")

from utils.torch_utils import select_device
from utils.augmentations import letterbox
import pycuda.autoinit
import pycuda.driver as cuda

try:
    import tensorrt as trt
except Exception as e:
    log.debug("importing tensorrt error: {}".format(str(e)))
import tensorrt as trt

if __name__ == "__main__":

    class time_count:
        def __init__(self):
            self.count = 0
            self.count1_call_times = 0
            self.names = []

        def init(self):
            self.count += 1
            self.count_call_times = 0

            if self.count == 1:
                pass
            elif self.count == 2:
                self.list = [[] for _ in range(self.count1_call_times)]
            elif self.count > 2:
                pass
            else:
                raise Exception("time_count")
            self.start = time.time()

        def __call__(self, name):
            self.end = time.time()
            if self.count == 1:
                self.count1_call_times += 1
                self.names.append(name)
            elif self.count > 1:
                self.list[self.count_call_times].append(self.end - self.start)
                self.count_call_times += 1
            else:
                raise Exception("time_count")
            self.start = time.time()

        def summury(self):
            print("==" * 20)
            spend_time_mean_overall = 0
            for idx in range(len(self.names)):
                name = self.names[idx]
                spend_time_list = self.list[idx]
                spend_time_mean = np.mean(spend_time_list) * 1000
                spend_time_std = np.std(spend_time_list) * 1000
                if spend_time_mean > 0.01:
                    print("{: <80}{: <20}{: <20}".format(name, round(spend_time_mean, 2), round(spend_time_std, 10)))
                spend_time_mean_overall += spend_time_mean
            print("{: <80}{: <20}".format("overall_time", round(spend_time_mean_overall, 1)))
            print("{: <80}{: <20}".format("overall_fps", round(1000 / spend_time_mean_overall, 1)))

            print("==" * 20)
else:
    class time_count:
        def __init__(self):
            pass

        def init(self):
            pass

        def __call__(self, name):
            pass

        def summury(self):
            pass

tc = time_count()


def get_img_path_batches(batch_size, img_dir):
    ret = []
    batch = []
    for root, dirs, files in os.walk(img_dir):
        for name in files:
            if len(batch) == batch_size:
                ret.append(batch)
                batch = []
            batch.append(os.path.join(root, name))
    if len(batch) > 0:
        ret.append(batch)
    return ret


def plot_one_box(x, img, color=None, label=None, line_thickness=None):
    """
    description: Plots one bounding box on image img,
                 this function comes from YoLov5 project.
    param:
        x:      a box likes [x1,y1,x2,y2]
        img:    a opencv image object
        color:  color to draw rectangle, such as (0,255,0)
        label:  str
        line_thickness: int
    return:
        no return

    """
    tl = (
            line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1
    )  # line/font thickness
    color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
        cv2.putText(
            img,
            label,
            (c1[0], c1[1] - 2),
            0,
            tl / 3,
            [225, 255, 255],
            thickness=tf,
            lineType=cv2.LINE_AA,
        )


class YoLov5TRT(object):

    def __init__(self, engine_file_path):
        # Create a Context on this device,
        # self.ctx = cuda.Device(0).make_context()
        stream = cuda.Stream()
        TRT_LOGGER = trt.Logger(trt.Logger.INFO)
        runtime = trt.Runtime(TRT_LOGGER)

        # Deserialize the engine from file
        with open(engine_file_path, "rb") as f:
            engine = runtime.deserialize_cuda_engine(f.read())
        context = engine.create_execution_context()

        host_inputs = []
        cuda_inputs = []
        host_outputs = []
        cuda_outputs = []
        bindings = []

        for binding in engine:
            print('bingding:', binding, engine.get_binding_shape(binding))
            size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
            dtype = trt.nptype(engine.get_binding_dtype(binding))
            # Allocate host and device buffers
            host_mem = cuda.pagelocked_empty(size, dtype)
            cuda_mem = cuda.mem_alloc(host_mem.nbytes)
            # Append the device buffer to device bindings.
            bindings.append(int(cuda_mem))
            # Append to the appropriate list.
            if engine.binding_is_input(binding):
                # self.input_w = engine.get_binding_shape(binding)[-1]
                # self.input_h = engine.get_binding_shape(binding)[-2]
                self.input_h = self.input_w = model_imgsz
                host_inputs.append(host_mem)
                cuda_inputs.append(cuda_mem)
            else:
                host_outputs.append(host_mem)
                cuda_outputs.append(cuda_mem)

        # Store
        self.stream = stream
        self.context = context
        self.engine = engine
        self.host_inputs = host_inputs
        self.cuda_inputs = cuda_inputs
        self.host_outputs = host_outputs
        self.cuda_outputs = cuda_outputs
        self.bindings = bindings
        self.batch_size = 1

    @staticmethod
    def post_process_convert(objects):
        # res = {
        #     "algorithm_data": {
        #         "is_alert": len(objects)>0,
        #         "target_count": len(objects),
        #         "target_info": objects
        #     },
        #     "model_data": {
        #         "objects": objects
        #     }
        # }
        slagcar = []
        for o in objects:
            if o["name"] == "slagcar":
                slagcar.append(o)
        res = {
            "algorithm_data": {
                "is_alert": len(slagcar) > 0,
                "target_count": len(slagcar),
                "target_info": slagcar
            },
            "model_data": {
                "objects": objects
            }
        }
        return json.dumps(res)

    def myinfer(self, a_raw_image):
        # threading.Thread.__init__(self)
        if a_raw_image is None:
            return self.post_process_convert([])
        tc.init()
        # self.ctx.push()
        # Restore
        stream = self.stream
        context = self.context
        engine = self.engine
        host_inputs = self.host_inputs
        cuda_inputs = self.cuda_inputs
        host_outputs = self.host_outputs
        cuda_outputs = self.cuda_outputs
        bindings = self.bindings

        batch_input_image = np.empty(shape=[self.batch_size, 3, self.input_h, self.input_w])
        input_image, image_raw, origin_h, origin_w = self.preprocess_image(a_raw_image)
        tc("preprocess_image")

        input_image = np.ascontiguousarray(input_image)  # 节约时间
        # Copy input image to host buffer
        np.copyto(host_inputs[0], input_image.ravel())
        tc("copyto")

        # Transfer input data to the GPU.
        cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)
        # cuda.memcpy_htod_async(cuda_inputs[0], input_image, stream)

        # Run inference.
        context.execute_async(batch_size=self.batch_size, bindings=bindings, stream_handle=stream.handle)
        # Transfer predictions back from the GPU.
        cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream)

        # Synchronize the stream
        stream.synchronize()
        tc("synchronize")
        # self.ctx.pop()
        output = host_outputs[0]

        result_boxes, result_scores, result_classid = self.post_process(output[0:6001], origin_h, origin_w)
        tc("post_process")

        detect_objs = []
        for j in range(len(result_boxes)):
            box = result_boxes[j]
            x0, y0, x1, y1 = box
            conf = result_scores[j]
            detect_objs.append({
                'x': int(x0),
                'y': int(y0),
                'width': int(x1 - x0),
                'height': int(y1 - y0),
                'confidence': float(conf),
                'name': categories[int(result_classid[j])]
            })

        returnthing = self.post_process_convert(detect_objs)

        return returnthing

    def destroy(self):
        # Remove any context from the top of the context stack, deactivating it.
        self.ctx.pop()

    def preprocess_image(self, raw_bgr_image):
        """
        description: Convert BGR image to RGB,
                     resize and pad it to target size, normalize to [0,1],
                     transform to NCHW format.
        param:
            input_image_path: str, image path
        return:
            image:  the processed image
            image_raw: the original image
            h: original height
            w: original width
        """
        stride = model_imgsz
        h, w, c = raw_bgr_image.shape
        # print(">> raw_bgr_image", raw_bgr_image.dtype)
        img = letterbox(raw_bgr_image, model_imgsz, stride=stride, auto=True)[0]  # BGR
        tc('letterbox')
        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        img = np.expand_dims(img, axis=0)
        img = np.ascontiguousarray(img)
        # img = torch.from_numpy(img.copy()).cuda()
        img = img / 255.0
        tc('255')

        return img, raw_bgr_image, h, w

    def xywh2xyxy(self, origin_h, origin_w, x):
        """
        description:    Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
        param:
            origin_h:   height of original image
            origin_w:   width of original image
            x:          A boxes numpy, each row is a box [center_x, center_y, w, h]
        return:
            y:          A boxes numpy, each row is a box [x1, y1, x2, y2]
        """
        y = np.zeros_like(x)
        r_w = self.input_w / origin_w
        r_h = self.input_h / origin_h
        if r_h > r_w:
            y[:, 0] = x[:, 0] - x[:, 2] / 2
            y[:, 2] = x[:, 0] + x[:, 2] / 2
            y[:, 1] = x[:, 1] - x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2
            y[:, 3] = x[:, 1] + x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2
            y /= r_w
        else:
            y[:, 0] = x[:, 0] - x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2
            y[:, 2] = x[:, 0] + x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2
            y[:, 1] = x[:, 1] - x[:, 3] / 2
            y[:, 3] = x[:, 1] + x[:, 3] / 2
            y /= r_h

        return y

    def post_process(self, output, origin_h, origin_w):
        """
        description: postprocess the prediction
        param:
            output:     A numpy likes [num_boxes,cx,cy,w,h,conf,cls_id, cx,cy,w,h,conf,cls_id, ...]
            origin_h:   height of original image
            origin_w:   width of original image
        return:
            result_boxes: finally boxes, a boxes numpy, each row is a box [x1, y1, x2, y2]
            result_scores: finally scores, a numpy, each element is the score correspoing to box
            result_classid: finally classid, a numpy, each element is the classid correspoing to box
        """
        # Get the num of boxes detected
        num = int(output[0])
        # Reshape to a two dimentional ndarray
        pred = np.reshape(output[1:], (-1, 6))[:num, :]
        # Do nms
        boxes = self.non_max_suppression(pred, origin_h, origin_w, conf_thres=CONF_THRESH, nms_thres=IOU_THRESHOLD)
        result_boxes = boxes[:, :4] if len(boxes) else np.array([])
        result_scores = boxes[:, 4] if len(boxes) else np.array([])
        result_classid = boxes[:, 5] if len(boxes) else np.array([])
        return result_boxes, result_scores, result_classid

    def bbox_iou(self, box1, box2, x1y1x2y2=True):
        """
        description: compute the IoU of two bounding boxes
        param:
            box1: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h))
            box2: A box coordinate (can be (x1, y1, x2, y2) or (x, y, w, h))
            x1y1x2y2: select the coordinate format
        return:
            iou: computed iou
        """
        if not x1y1x2y2:
            # Transform from center and width to exact coordinates
            b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
            b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
            b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
            b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
        else:
            # Get the coordinates of bounding boxes
            b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
            b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]

        # Get the coordinates of the intersection rectangle
        inter_rect_x1 = np.maximum(b1_x1, b2_x1)
        inter_rect_y1 = np.maximum(b1_y1, b2_y1)
        inter_rect_x2 = np.minimum(b1_x2, b2_x2)
        inter_rect_y2 = np.minimum(b1_y2, b2_y2)
        # Intersection area
        inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, None) * \
                     np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, None)
        # Union Area
        b1_area = (b1_x2 - b1_x1 + 1) * (b1_y2 - b1_y1 + 1)
        b2_area = (b2_x2 - b2_x1 + 1) * (b2_y2 - b2_y1 + 1)

        iou = inter_area / (b1_area + b2_area - inter_area + 1e-16)

        return iou

    def non_max_suppression(self, prediction, origin_h, origin_w, conf_thres=0.5, nms_thres=0.4):
        """
        description: Removes detections with lower object confidence score than 'conf_thres' and performs
        Non-Maximum Suppression to further filter detections.
        param:
            prediction: detections, (x1, y1, x2, y2, conf, cls_id)
            origin_h: original image height
            origin_w: original image width
            conf_thres: a confidence threshold to filter detections
            nms_thres: a iou threshold to filter detections
        return:
            boxes: output after nms with the shape (x1, y1, x2, y2, conf, cls_id)
        """
        # Get the boxes that score > CONF_THRESH
        boxes = prediction[prediction[:, 4] >= conf_thres]
        # Trandform bbox from [center_x, center_y, w, h] to [x1, y1, x2, y2]
        boxes[:, :4] = self.xywh2xyxy(origin_h, origin_w, boxes[:, :4])
        # clip the coordinates
        boxes[:, 0] = np.clip(boxes[:, 0], 0, origin_w - 1)
        boxes[:, 2] = np.clip(boxes[:, 2], 0, origin_w - 1)
        boxes[:, 1] = np.clip(boxes[:, 1], 0, origin_h - 1)
        boxes[:, 3] = np.clip(boxes[:, 3], 0, origin_h - 1)
        # Object confidence
        confs = boxes[:, 4]
        # Sort by the confs
        boxes = boxes[np.argsort(-confs)]
        # Perform non-maximum suppression
        keep_boxes = []
        while boxes.shape[0]:
            large_overlap = self.bbox_iou(np.expand_dims(boxes[0, :4], 0), boxes[:, :4]) > nms_thres
            label_match = boxes[0, -1] == boxes[:, -1]
            # Indices of boxes with lower confidence scores, large IOUs and matching labels
            invalid = large_overlap & label_match
            keep_boxes += [boxes[0]]
            boxes = boxes[~invalid]
        boxes = np.stack(keep_boxes, 0) if len(keep_boxes) else np.array([])
        return boxes


def pre_convert():
    # log.info("pre_convert ")
    pt_file = "/project/train/models/yolov5/{}/weights/{}".format(model_name, pt_name)
    wts_file = "/project/train/models/yolov5/{}/weights/{}.wts".format(model_name, model_name)
    engine_file_path = "/project/train/models/yolov5/{}/weights/{}.engine".format(model_name, model_name)

    if not os.path.exists(wts_file):
        # Initialize
        device = select_device('cuda:0')
        model = torch.load(pt_file, map_location=device)['model'].float()  # load to FP32
        delattr(model.model[-1], 'anchor_grid')  # model.model[-1] is detect layer
        # The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight.
        model.model[-1].register_buffer("anchor_grid", torch.Tensor(model.yaml['anchors']))
        model.to(device).eval()
        with open(wts_file, 'w') as f:
            f.write('{}\n'.format(len(model.state_dict().keys())))
            for k, v in model.state_dict().items():
                vr = v.reshape(-1).cpu().numpy()
                f.write('{} {} '.format(k, len(vr)))
                for vv in vr:
                    f.write(' ')
                    f.write(struct.pack('>f', float(vv)).hex())
                f.write('\n')
        log.info("converted wts " + wts_file + "\n\n")

    if not os.path.exists(engine_file_path):
        # log.info("converting engine " + engine_file_path + "\n\n")
        os.system(
            "cd /project/tensorrtx/yolov5/{}; sudo /project/tensorrtx/yolov5/{}/yolov5 -s  {}  {}  {}".format(bbuild,
                                                                                                              bbuild,
                                                                                                              wts_file,
                                                                                                              engine_file_path,
                                                                                                              model_type))
        # log.info("converted engine " + engine_file_path + "\n\n")
    # if not os.path.exists(engine_file_path):
    #     log.info("engine_file_path failed " + engine_file_path)
    #     raise Exception("engine_file_path failed " + engine_file_path)


def init():
    pre_convert()

    # load custom plugin and engine
    PLUGIN_LIBRARY = "/project/tensorrtx/yolov5/build/libmyplugins.so"
    engine_file_path = "/project/train/models/yolov5/{}/weights/{}.engine".format(model_name, model_name)
    ctypes.CDLL(PLUGIN_LIBRARY)

    # a YoLov5TRT instance
    yolov5_wrapper = YoLov5TRT(engine_file_path)
    # yolov5_wrapper.myinfer(np.zeros((model_imgsz, model_imgsz, 3), dtype=np.uint8))

    return yolov5_wrapper


def process_image(yolov5_wrapper, input_image, *args, **kwargs):
    res = yolov5_wrapper.myinfer(input_image)

    return res


def show(img, result, show_path):
    result = json.loads(result)
    detect_objs = result["objects"]
    imgshow = img.copy()
    if len(detect_objs) > 0:
        for ann in detect_objs:
            xmin = int(ann["xmin"])
            ymin = int(ann["ymin"])
            xmax = int(ann["xmax"])
            ymax = int(ann["ymax"])
            confidence = float(ann["confidence"])
            name = ann["name"]
            imgshow = cv2.rectangle(imgshow, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2)
            imgshow = cv2.putText(imgshow, name, (int(xmin), int(ymin)), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 2)
            imgshow = cv2.putText(imgshow, str(int(confidence * 100)), (int(xmin), int(ymin)), cv2.FONT_HERSHEY_SIMPLEX,
                                  1.2, (0, 0, 255), 2)
    cv2.imwrite(show_path, imgshow)


if __name__ == '__main__':

    imgpaths = glob.glob("/home/data/newdata/*.jpg")
    # imgpath = imgpaths[int((len(imgpaths)-1)*random.random())]
    # show_dir = "/project/train/show_ji"
    # if not os.path.exists(show_dir):
    #     os.mkdir(show_dir)
    # show_path = show_dir+ "/" + os.path.basename(imgpath)
    # ============================================================================
    predictor = init()
    for imgpath in tqdm(imgpaths):
        img = cv2.imread(imgpath)
        result = process_image(predictor, img)
        # print(result)
        # show(img, result, show_path)
    # yolov5_wrapper.destroy()
    tc.summury()

